import torch
from PIL import Image
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
import json
import pandas as pd

device = "cuda"
model_path = "/rainbow/mnt/ai_pro/sd/swift/glm_4v_9b"


model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    trust_remote_code=True
).to(device).eval()


prompt_txt = f"""
如图是一个购物小票，图中是否出现 “AI助手” 字样，如果有输出 “存在”，如果没有则输出 “不存在”。
"""

print(prompt_txt)


def model_result(img_path):
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    query = prompt_txt
    image = Image.open(img_path).convert('RGB')
    inputs = tokenizer.apply_chat_template([{"role": "user", "image": image, "content": query}],
                                           add_generation_prompt=True, tokenize=True, return_tensors="pt",
                                           return_dict=True)  # chat mode
    
    inputs = inputs.to(device)
    
    gen_kwargs = {"max_length": 2500, "do_sample": True, "top_k": 1}
    with torch.no_grad():
        outputs = model.generate(**inputs, **gen_kwargs)
        outputs = outputs[:, inputs['input_ids'].shape[1]:]
        result_=tokenizer.decode(outputs[0])
    return result_



files = os.listdir("img_dir/")
print(files)
for file in files:
    file_path = os.path.join("/rainbow/mnt/ai_pro/sd/glm_4v_9b/img_dir", file)
    print("file_path:",file_path)
    try:
        result_=model_result(file_path)
        print('model_result:',result_)
    except Exception as e:
        print('find {  } 错误:',e)

"""


df1=pd.DataFrame(columns=["v1_商户名称","v2_订单编号","v3_交易日期","v4_交易金额"])
files = os.listdir("img_dir/")
print(files)
for file in files:
    file_path = os.path.join("/rainbow/mnt/ai_pro/sd/glm_4v_9b/img_dir", file)
    print("file_path:",file_path)

    result_=str('''{"商户名称": "-","订单编号": "-","交易日期": "-","交易金额": "-"}''')
    try:
        result_=model_result(file_path)
        #print('model_result:',result_)
        index_left = result_.find('{') 
        index_right = result_.find('}')
        result_=result_[index_left:index_right+1]
    except Exception as e:
        print('find {  } 错误:',e)
    try:
        json_data=json.loads(result_)
    except Exception as e:
        print('json.loads 错误:',e) 
        json_data=json.loads('{"商户名称": "-","订单编号": "-","交易日期": "-","交易金额": "-"}')
             
    print("aaaaaaaa::\n",type(json_data),json_data)

    df1.loc[files.index(file),'v1_商户名称']=json_data["商户名称"]
    df1.loc[files.index(file),'v2_订单编号']=str(json_data["订单编号"])
    df1.loc[files.index(file),'v3_交易日期']=str(json_data["交易日期"])
    df1.loc[files.index(file),'v4_交易金额']=str(json_data["交易金额"])

df1.to_excel("result_glm_test.xlsx",index=False)
"""